# -*- coding: utf-8 -*- #

# -----------------------------------------------------------------------
# File Name:    inference.py
# Version:      ver1_0
# Created:      2024/06/17
# Description:  本文件定义了用于在模型应用端进行推理，返回模型输出的流程
#               ★★★请在空白处填写适当的语句，将模型推理应用流程补充完整★★★
# -----------------------------------------------------------------------

import torch
from PIL import Image
from torchvision.transforms import ToTensor


def inference(image_path, model, device):
    """定义模型推理应用的流程。"""
    # 将模型置为评估（测试）模式
    model.eval()

    try:
        # 加载并预处理图像
        image = Image.open(image_path)
        transform = ToTensor()
        image_tensor = transform(image).unsqueeze(0).to(device)  # 增加batch维度

        # 进行推理
        with torch.no_grad():
            output = model(image_tensor)
            _, predicted = torch.max(output.data, 1)
            prediction = predicted.item()

        # 显示结果
        image.show()
        print(f"预测类别: {prediction}")

        return prediction

    except Exception as e:
        print(f"推理过程中发生错误: {str(e)}")
        return None


if __name__ == "__main__":
    # 指定图片路径
    image_path = "./images/test/signs/img_0001.png"

    # 加载训练好的模型
    model = torch.load('./models/model.pkl')
    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")
    model.to(device)

    # 显示图片，输出预测结果
    inference(image_path, model, device)
